#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : Exchange.py
# @Author: Richard Chiming Xu
# @Date  : 2022/1/25
# @Desc  : 交换数据增强

'''
    数据集中每条数据有两个句子
    a, b, 1
    a, c, 1
    a, d, 0
    a~b, a~c => b~c
    a~b, ad不相似 => bd不相似
'''

import pandas as pd
from collections import defaultdict


def aug_group_by_a(df):
    aug_data = defaultdict(list)
    # 以text_a中的句子为 a
    for g, data in df.groupby(by=['text_a']):
        if len(data) < 2:
            continue
        for i in range(len(data)):
            for j in range(i + 1, len(data)):
                # 取出b的值，a,b的label
                row_i_text = data.iloc[i, 1]
                row_i_label = data.iloc[i, 2]

                # 取出c的值，a,c的label
                row_j_text = data.iloc[j, 1]
                row_j_label = data.iloc[j, 2]

                if row_i_label == row_j_label == 0:
                    continue

                aug_label = 1 if row_i_label == row_j_label == 1 else 0

                aug_data['text_a'].append(row_i_text)
                aug_data['text_b'].append(row_j_text)
                aug_data['label'].append(aug_label)
    return pd.DataFrame(aug_data)

if __name__ == '__main__':
    origin_df = pd.read_csv('../data/paws-x/train.tsv', sep='\t', names=['text_a', 'text_b', 'label'])
    aug_df = aug_group_by_a(origin_df)
    print(aug_df[aug_df['label'] == 1])
